# -*-Python-*-

import mesh_tensorflow.optimize
import mesh_tensorflow.transformer.learning_rate_schedules
import mesh_tensorflow.transformer.transformer
import mesh_tensorflow.transformer.transformer_layers
import mesh_tensorflow.transformer.learning_rate_schedules

# for Unitransformer models (e.g. language-model)
transformer.make_layer_stack.layers = [
    @mesh_tensorflow.transformer.transformer_layers.SelfAttention,
    @mesh_tensorflow.transformer.transformer_layers.DenseReluDense,
]

# for Bitransformer models (two-stack sequence-to-sequence models)
encoder/transformer.make_layer_stack.layers = [
    @mesh_tensorflow.transformer.transformer_layers.SelfAttention,
    @mesh_tensorflow.transformer.transformer_layers.DenseReluDense,
]
decoder/transformer.make_layer_stack.layers = [
    @mesh_tensorflow.transformer.transformer_layers.SelfAttention,
    @mesh_tensorflow.transformer.transformer_layers.EncDecAttention,
    @mesh_tensorflow.transformer.transformer_layers.DenseReluDense,
]

transformer.LayerStack.sublayers_initial = [
    @transformer.sublayer_dropout,
]
transformer.LayerStack.sublayers_per_layer = [
    @transformer.sublayer_rms_norm,
    @transformer.sublayer_call_layer,
    @transformer.sublayer_dropout,
    @transformer.sublayer_residual,
]
transformer.LayerStack.sublayers_final = [
    @transformer.sublayer_rms_norm,
    @transformer.sublayer_dropout,
]

Unitransformer.shared_embedding_and_softmax_weights = True

# Model sizes
num_layers = 6
d_model = 1024
d_ff = 4096
num_heads = 8
d_kv = 128

# These are deprecated hyperparameters that trigger legacy behavior
# for compatibility with old checkpoints.
LayerStack.dropout_rate = None
LayerStack.norm_epsilon = None

transformer.make_layer_stack.num_layers = %num_layers
Unitransformer.d_model = %d_model
DenseReluDense.hidden_size = %d_ff
SelfAttention.num_heads = %num_heads
SelfAttention.key_value_size = %d_kv

# dropout
dropout_rate = 0.1
DenseReluDense.dropout_rate = %dropout_rate
SelfAttention.dropout_rate = %dropout_rate
LocalSelfAttention.dropout_rate = %dropout_rate
TalkingHeadsSelfAttention.dropout_rate = %dropout_rate
GeneralBilinearSelfAttention.dropout_rate = %dropout_rate
transformer.sublayer_dropout.dropout_rate = %dropout_rate

# No label smoothing by default
decoder/Unitransformer.label_smoothing = 0.0

# the sequence length for this run
utils.run.sequence_length = {"inputs": 512, "targets": 512}

# The size of the positional embedding tables.
# This is relevant only if positional embeddings are used.
# The value must be at least as large as the actual sequence length.
# Changing this value breaks checkpoint compatibility.
Unitransformer.max_length = 512

# Optimization
utils.run.optimizer = @optimize.AdafactorOptimizer

# The following layout rules have no effect unless the names "model" and/or
# "batch" are in utils.run.mesh_shape .
utils.run.layout_rules = "ensemble:ensemble,batch:batch,d_ff:model,heads:model,vocab:model,experts:batch"

# During training, serialize large batches to save memory.
utils.serialize_num_microbatches.tokens_per_microbatch_per_replica = 2048

# Learning rate schedule including rqrt-decay followed by linear-decay
utils.run.learning_rate_schedule = @learning_rate_schedules.learning_rate_schedule_noam

# Default to data-parallelism
utils.run.mesh_shape = @mesh_tensorflow.transformer.utils.tpu_mesh_shape()

# Override this for model-parallelism
utils.tpu_mesh_shape.model_parallelism = 1

# bfloat16 activations for faster allreduce
# Note that bfloat16 kernels do not exist for CPU, so you will need to remove
# this if you want to run on CPU. It can also cause some numerical
# instabilities, but it is often significantly faster.
utils.get_variable_dtype.activation_dtype = "bfloat16"
